#!/usr/bin/env python3
"""
最简 RAG 单次补丁生成器（OpenAI 线上模型）
运行后得到 preds.json → 直接 sb-cli 提交
"""
import json, os, glob, tiktoken, tqdm, argparse
from datasets import load_dataset
from openai import OpenAI

client = OpenAI(
    base_url="http://115.190.56.170:8000/v1",  # 注意最后不要漏掉 /v1
    api_key="sk-4duyWc4ZphbqBIAQytlp3A"
)

# ---------- 参数区 ----------
MODEL = "qwen3-coder-480b-a35b-instruct"  # 也可 gpt-4 / gpt-4-turbo
MAX_CODE_TOKENS = 3000  # 控制代码长度，防止超长
MAX_OUTPUT = 1024  # 最大 diff token
TEMPERATURE = 0.5


def call_openai(prompt: str) -> str:
    """

    :rtype: str
    """
    rsp = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=TEMPERATURE,
        max_tokens=MAX_CODE_TOKENS,
    )
    return rsp.choices[0].message.content or ""

def main():
    prompt=f"你好，你是谁？"
    raw = call_openai(prompt)
    print(raw)

if __name__ == '__main__':
    main()

